package com.terry.trace.plugin;

import com.netflix.hystrix.HystrixThreadPoolKey;
import com.netflix.hystrix.HystrixThreadPoolProperties;
import com.netflix.hystrix.strategy.HystrixPlugins;
import com.netflix.hystrix.strategy.concurrency.HystrixConcurrencyStrategy;
import com.netflix.hystrix.strategy.concurrency.HystrixRequestVariable;
import com.netflix.hystrix.strategy.concurrency.HystrixRequestVariableLifecycle;
import com.netflix.hystrix.strategy.eventnotifier.HystrixEventNotifier;
import com.netflix.hystrix.strategy.executionhook.HystrixCommandExecutionHook;
import com.netflix.hystrix.strategy.metrics.HystrixMetricsPublisher;
import com.netflix.hystrix.strategy.properties.HystrixPropertiesStrategy;
import com.netflix.hystrix.strategy.properties.HystrixProperty;
import com.terry.trace.core.LogTagContext;
import com.terry.trace.core.TraceLogger;

import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * hystrix并发策略，用来传递logTag到hystrix的线程池
 * @author TerryJiangConcurrencyStrategy
 */
public class ThreadLocalCallableHystrixConcurrencyStrategy extends HystrixConcurrencyStrategy {

	private static final TraceLogger LOGGER = TraceLogger.getLogger(ThreadLocalCallableHystrixConcurrencyStrategy.class);

	/** 代理的已存在的hystrix并发策略 */
	private HystrixConcurrencyStrategy delegate;

	/**
	 * 构建并发策略
	 */
	public ThreadLocalCallableHystrixConcurrencyStrategy() {
		try {
			// 获取已存在的并发策略
			this.delegate = HystrixPlugins.getInstance().getConcurrencyStrategy();
			// 如果是已经注册的是当前相同策略，则无需再注册
			if (this.delegate instanceof ThreadLocalCallableHystrixConcurrencyStrategy) {
				// Welcome to singleton hell...
				return;
			}

			// 获取代理的并发策略相关的其他策略
			HystrixCommandExecutionHook commandExecutionHook = HystrixPlugins.getInstance().getCommandExecutionHook();
			HystrixEventNotifier eventNotifier = HystrixPlugins.getInstance().getEventNotifier();
			HystrixMetricsPublisher metricsPublisher = HystrixPlugins.getInstance().getMetricsPublisher();
			HystrixPropertiesStrategy propertiesStrategy = HystrixPlugins.getInstance().getPropertiesStrategy();

			this.logCurrentStateOfHystrixPlugins(eventNotifier, metricsPublisher, propertiesStrategy);

			// 重置hystrix所有插件。因为一类插件只能注册一个实例，如果已经注册了，则无法再注册其他实例，所以需要重置
			HystrixPlugins.reset();

			// 注册当前并发策略
			HystrixPlugins.getInstance().registerConcurrencyStrategy(this);
			HystrixPlugins.getInstance().registerCommandExecutionHook(commandExecutionHook);
			HystrixPlugins.getInstance().registerEventNotifier(eventNotifier);
			HystrixPlugins.getInstance().registerMetricsPublisher(metricsPublisher);
			HystrixPlugins.getInstance().registerPropertiesStrategy(propertiesStrategy);
		}
		catch (Exception e) {
			LOGGER.error("Failed to register ThreadLocal Hystrix Concurrency Strategy", e);
		}
	}

	/**
	 * 打印hystrix插件的当前状态
	 * @param eventNotifier
	 * @param metricsPublisher
	 * @param propertiesStrategy
	 */
	private void logCurrentStateOfHystrixPlugins(HystrixEventNotifier eventNotifier, HystrixMetricsPublisher metricsPublisher,
												 HystrixPropertiesStrategy propertiesStrategy) {
		LOGGER.info("Current Hystrix plugins configuration is ["
				+ "concurrencyStrategy [" + this.delegate + "]," + "eventNotifier ["
				+ eventNotifier + "]," + "metricPublisher [" + metricsPublisher + "],"
				+ "propertiesStrategy [" + propertiesStrategy + "]," + "]");
		LOGGER.info("Registering ThreadLocal Hystrix Concurrency Strategy.");
	}

	/**
	 * 包装callable线程任务
	 * @param callable
	 * @param <T>
	 * @return
	 */
	@Override
	public <T> Callable<T> wrapCallable(Callable<T> callable) {
		// 如果任务已经是当前自定义的任务，无需再进行包装
		if (callable instanceof ThreadLocalCallableHystrixConcurrencyStrategy.TraceAwareCallable) {
			return callable;
		}

		// 如果代理的线程任务不为空，则说明有代理任务。有代理任务，则代理任务也进行包装
		Callable<T> wrappedCallable = this.delegate != null ? this.delegate.wrapCallable(callable) : callable;

		// 如果包装后的任务是当前自定义的任务，无需再进行包装
		if (wrappedCallable instanceof ThreadLocalCallableHystrixConcurrencyStrategy.TraceAwareCallable) {
			return wrappedCallable;
		}

		// 包装成当前自定义的任务
		return new ThreadLocalCallableHystrixConcurrencyStrategy.TraceAwareCallable<>(callable, LogTagContext.get());
	}

	/**
	 * trace透传线程任务
	 * @param <T>
	 */
	static class TraceAwareCallable<T> implements Callable<T> {

		/** 代理任务 */
		private final Callable<T> delegate;
		/** log tag */
		private final Map<String, String> logTagMap;

		/**
		 * 通过构造器传递父线程的log tag, 并代理原有的线程任务
		 * @param callable
		 * @param logTagMap
		 */
		public TraceAwareCallable(Callable<T> callable, Map<String, String> logTagMap) {
			this.delegate = callable;
			this.logTagMap = logTagMap;
		}

		@Override
		public T call() throws Exception {
			try {
				// 将父线程中的log tag设置到本地线程上下文
				LogTagContext.set(logTagMap);
				// 执行代理线程任务
				return delegate.call();
			} finally {
				// 要放在最后一步执行，关键步骤，清空上下文中的log tag, 用来防止内存泄露，或者产生脏数据
				LogTagContext.remove();
			}
		}
	}

	/**
	 * 获取线程池（通过具体的线程池参数构造线程池）
	 * 使用代理的并发策略的线程池策略，因为没有特殊要求，且不影响到被代理策略
	 * @param threadPoolKey
	 * @param corePoolSize
	 * @param maximumPoolSize
	 * @param keepAliveTime
	 * @param unit
	 * @param workQueue
	 * @return
	 */
	@Override
	public ThreadPoolExecutor getThreadPool(HystrixThreadPoolKey threadPoolKey,
											HystrixProperty<Integer> corePoolSize,
											HystrixProperty<Integer> maximumPoolSize,
											HystrixProperty<Integer> keepAliveTime, TimeUnit unit,
											BlockingQueue<Runnable> workQueue) {
		return this.delegate.getThreadPool(threadPoolKey, corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
	}

	/**
	 * 获取线程池(通过threadPoolProperties线程属性类构造线程池)
	 * 使用代理的并发策略的线程池策略，因为没有特殊要求，且不影响到被代理策略
	 * @param threadPoolKey
	 * @param threadPoolProperties
	 * @return
	 */
	@Override
	public ThreadPoolExecutor getThreadPool(HystrixThreadPoolKey threadPoolKey, HystrixThreadPoolProperties threadPoolProperties) {
		return this.delegate.getThreadPool(threadPoolKey, threadPoolProperties);
	}

	/**
	 * 获取blocking queue
	 * 使用代理的并发策略的获取blocking queue的方式，因为没有特殊要求，且不影响到被代理策略
	 * @param maxQueueSize
	 * @return
	 */
	@Override
	public BlockingQueue<Runnable> getBlockingQueue(int maxQueueSize) {
		return this.delegate.getBlockingQueue(maxQueueSize);
	}

	/**
	 * 获取请求变量
	 * 使用代理的并发策略的获取请求变量的方式，因为没有特殊要求，且不影响到被代理策略
	 * @param rv
	 * @param <T>
	 * @return
	 */
	@Override
	public <T> HystrixRequestVariable<T> getRequestVariable(HystrixRequestVariableLifecycle<T> rv) {
		return this.delegate.getRequestVariable(rv);
	}
}
